function updated_posterior = optimize_rsl(posterior, prior, data, n_iterations);
%
% Learns an updated posterior approximation for parameters of the Relevant Subtask
% Learning (RSL) model, by iterative optimization of the approximation, based on 
% the data, a prior, and an initial posterior approximation.
%
% Inputs: posterior = initial posterior approximation, given by the initialize_rsl
% function. prior = prior for parameters, also given by the initialize_rsl function.
% data = cell array of the several tasks, in the same format as for initialize_rsl
% (first task is the task of interest). n_iterations: number of iterations used in 
% the optimization; 500 is a good number of iterations.
%
% Outputs: updated posterior approximation.
%


start_time = clock;

n_tasks = length(data);


Sigma_U0 = prior{1};
theta_U0 = prior{2};


for iteration=1:n_iterations,
  previous_posterior = posterior;
  
  %-----------------------------------------
  % Part 1: update approximation parameters
  %-----------------------------------------
  
  % update xis_S_U for task-of-interest
  Sigma_U = posterior{1};
  theta_U = posterior{2};  
  X = data{1};
  cs = X(:,end)';
  X = X(:,1:end-1)';  
  xis_U = sum(X.*((Sigma_U+theta_U*theta_U')*X),1).^0.5;
  
  posterior{3} = xis_U;
  
  for S=2:n_tasks,
    X = data{S};
    cs = X(:,end)';
    X = X(:,1:end-1)';
    n_S = size(X,2);
    n_D = size(X,1);    
    Sigma_S = posterior{(S-2)*7+4};
    theta_S = posterior{(S-2)*7+5};
    
    % update xis_S
    xis_S = sum(X.*((Sigma_S+theta_S*theta_S')*X),1).^0.5;
    posterior{(S-2)*7+9} = xis_S;
    
    % update xis_U
    xis_U = sum(X.*((Sigma_U+theta_U*theta_U')*X),1).^0.5;
    posterior{(S-2)*7+10} = xis_U;        
  end;
  
  
  %------------------------------------------
  % Part 2: update model parameters
  %------------------------------------------  
  
  % update Sigma_U and theta_U
  Sigma_U = inv(Sigma_U0);
  theta_U = inv(Sigma_U0)*theta_U0;

  for S=1:n_tasks,
    X = data{S};
    cs = X(:,end)';
    X = X(:,1:end-1)';
    n_S = size(X,2);
    n_D = size(X,1);

    
    if S == 1,
      gammas = zeros(1,n_S);
      xis_U = posterior{3};
    else
      gammas = posterior{(S-2)*7+8};      
      xis_U = posterior{(S-2)*7+10};
    end;  
        
    % compute weighing function g(\xi_{S,i,U})
    I = find(xis_U == 0);
    gs_U = zeros(1,n_S);
    gs_U(I) = 1;
    I2 = find(xis_U ~= 0);
    gs_U(I2) = exp(-xis_U(I2));
    gs_U(I2) = 2*(1-gs_U(I2))./(xis_U(I2).*(1+gs_U(I2)));
   
    for i=1:n_S,    
      Sigma_U = Sigma_U + 0.25*(1-gammas(i))*gs_U(i)*X(:,i)*X(:,i)';
      theta_U = theta_U + 0.5*(1-gammas(i))*cs(i)*X(:,i);  
    end;

  end;
  Sigma_U = inv(Sigma_U+1e-6*eye(n_D));
  posterior{1} = Sigma_U;

  theta_U = Sigma_U*theta_U;
  posterior{2} = theta_U;
 
  
  
  % update posteriors for task-specific parameters
  for S=2:n_tasks,
    X = data{S};
    cs = X(:,end)';
    X = X(:,1:end-1)';
    n_S = size(X,2);
    n_D = size(X,1);
    
    Sigma_S0 = prior{(S-2)*4+3};
    theta_S0 = prior{(S-2)*4+4};
    phi_SA0 = prior{(S-2)*4+5};
    phi_SB0 = prior{(S-2)*4+6};
    
    Sigma_S = posterior{(S-2)*7+4};
    theta_S = posterior{(S-2)*7+5};
    phi_SA = posterior{(S-2)*7+6};
    phi_SB = posterior{(S-2)*7+7};
    gammas = posterior{(S-2)*7+8};
    xis_S = posterior{(S-2)*7+9};
    xis_U = posterior{(S-2)*7+10};

    % compute weighing function g(\xi_{S,i,S})
    gs_S = zeros(1,n_S);
    I = find(xis_S == 0);
    gs_S(I) = 1;
    I2 = find(xis_S ~= 0);
    gs_S(I2) = exp(-xis_S(I2));
    gs_S(I2) = 2*(1-gs_S(I2))./(xis_S(I2).*(1+gs_S(I2)));
    
    % compute weighing function g(\xi_{S,i,U})
    gs_U = zeros(1,n_S);
    I = find(xis_U == 0);
    gs_U(I) = 1;
    I2 = find(xis_U ~= 0);
    gs_U(I2) = exp(-xis_U(I2));
    gs_U(I2) = 2*(1-gs_U(I2))./(xis_U(I2).*(1+gs_U(I2)));
    
    
    % update Sigma_S

    Sigma_S = inv(Sigma_S0);    
    theta_S = inv(Sigma_S0)*theta_S0;
    for i=1:n_S,
      Sigma_S = Sigma_S + 0.25*gammas(i)*gs_S(i)*X(:,i)*X(:,i)';
      theta_S = theta_S + 0.5*(gammas(i))*cs(i)*X(:,i); 
    end;
    Sigma_S = inv(Sigma_S+1e-6*eye(n_D));
    theta_S = Sigma_S*theta_S;
            
    posterior{(S-2)*7+4} = Sigma_S;   
    posterior{(S-2)*7+5} = theta_S;
    
    % update gammas
    hs1 = 0.5*cs.*(theta_S'*X) - 0.125*gs_S.*sum(X.*((Sigma_S+theta_S*theta_S')*X),1) - ...
          log(1+exp(-xis_S)) - 0.5*xis_S + 0.125*gs_S.*(xis_S.^2) + ...
          psi(phi_SB);
    
    hs2 = 0.5*cs.*(theta_U'*X) - 0.125*gs_U.*sum(X.*((Sigma_U+theta_U*theta_U')*X),1) - ...
          log(1+exp(-xis_U)) - 0.5*xis_U + 0.125*gs_U.*(xis_U.^2) + ...
          psi(phi_SA);
    C = max([max(hs1) max(hs2)]);
    gammas = exp(hs1-C)./(exp(hs1-C)+exp(hs2-C));
    posterior{(S-2)*7+8} = gammas;
    
    % update phi_SA
    phi_SA = phi_SA0 + sum(1-gammas);
    posterior{(S-2)*7+6} = phi_SA;
    
    % update phi_SB
    phi_SB = phi_SB0 + sum(gammas);
    posterior{(S-2)*7+7} = phi_SB;    
  end;

  if mod(iteration,10)==0,

    % compute difference between posteriors;
    postdiff=0;
    postdiff = postdiff + sum(sum((posterior{1}-previous_posterior{1}).^2));
    postdiff = postdiff + sum((posterior{2}-previous_posterior{2}).^2);

    current_time = clock;
    elapsed=etime(current_time,start_time)/60;
    fprintf(1, '\r iteration %d, postdiff %e (%.2f min spent, %.2f left)', ...
            iteration, postdiff, elapsed, ...
            (n_iterations-iteration)*elapsed/iteration);
  end;
end;

fprintf(1, '...done.\n');

updated_posterior = posterior;

